#!/usr/bin/env python

import numpy as np
import time
import sys
import os
import subprocess
from math import pi
from xyz_src.write_lammps import run_umb

import sys

def loop(init_config, job_index, batch, modeldir, workdir, r, k, ncpus, maxK, temp=100, damp=100):

    data = np.load(init_config)
    allz = data['z']
    minz = np.min(allz)
    maxz = np.max(allz)
    print(minz, maxz)

    case = len(r)-1

    ind = np.argmin(np.abs(r[job_index]-allz))
    # # ind = np.argmin(np.abs(allz))
    # if r[job_index] > 4:
    #     ind = np.argmax(allz)
    # elif r[job_index] < -4:
    #     ind = np.argmin(allz)
    # else:
    #     ind = np.argmin(np.abs(allz))
    print(data['xyz'].shape)
    print(data['z'].shape)
    x = data['xyz'][ind]
    z = allz[ind]
    print("initial config", r[job_index], z[-1])

    results = []
    whole_runlist = [i for i in range(case)] #[10, 20, 30, 40, 45, 50]
    cur_workdir = f"{workdir}_{r[job_index]}"
    if not os.path.exists(cur_workdir):
        os.mkdir(cur_workdir)
    print("compute", r[job_index], k[job_index])
    umb(x, r[job_index], k[job_index], r[job_index+1]-r[job_index],
            cur_workdir, modeldir, maxK, temp, damp)

def umb(x0, R0, K, dx, workdir, modeldir, maxK, temp, damp):

    stop = False
    idx = 0
    while (not stop):
        t0 = time.time()
        b, e, xi, maxe = run_umb(idx, x0, vxyzflat=None, temp=temp,
                         step=2000000, center=R0, k=K,
                         model=modeldir, stride=100,
                         workdir=workdir,
                         damp=damp)
        print("timer", time.time()-t0)
        idx += 1
        ave_xi = np.average(xi)
        if (np.abs(ave_xi-R0)<dx):
            stop = True
        if (K>maxK):
            stop = True
        if maxe > 10:
            stop = True
        K = K*2
